#!/usr/bin/python
# -*- coding: UTF-8 -*-
# author: ziping.yu
# Date: 2023.5.10
import rospy
import PyKDL
from geometry_msgs.msg import Twist, Point, PoseStamped
from std_msgs.msg import String
import numpy as np
import math
import sys

class NavigationPlan:
  def __init__(self, drone_type, drone_id, drone_num=0, vehicle_num=0):
    self.drone_type = drone_type
    self.drone_id = int(drone_id)
    self.drone_num = int(drone_num)
    self.vehicle_num = int(vehicle_num)
    self.bias = np.array([[0,0,0], [3,3,0], [3,-3,0]]) # only for uav which in enu coordinate
    
    '''
    ros parameter configuration
    '''
    self.rate = 30 # pin lv
    
    '''
    control parameter configuration
    '''
    self.kp = 0.5
    self.kv = 1.0
    self.ky = 1.0
    
    '''
    boundary value configuration
    '''
    self.vxy_max = 4
    self.vz_max = 3
    self.angz_max = 3
    self.height_limit = 5.0
    self.safe_distance = 0.5 # distance between drones for avoidance
    self.safe_height = 0.5
    '''
    variable of drone
    '''
    self.current_pose = PoseStamped()
    self.current_position = Point()
    self.current_yaw = 0.0
    self.current_vel = Twist()
    self.target_motion = Twist()
    self.target_position = Twist()
    self.target_yaw = 0.0
    self.last_yaw = 0.0
    self.cmd = ''
    self.other_drone_pose = [PoseStamped() for i in range(self.drone_num)]
    self.other_vehicle_pose = [PoseStamped() for i in range(self.vehicle_num)]
    
    '''
    FSM variable
    '''
    self.FSM_EXEC_STATE = {0:"TAKEOFF", 1:"INIT", 2:"GEN_TRAJ", 3:"EXEC"}
    self.exec_state = 0
    self.change_task_flag = False
    self.enable_receive_goal = False
    self.count = 0
    self.is_arrive = False
    self.arrive_count = 0
    
    rospy.init_node(f"{drone_type}_{str(drone_id)}_plan")
    '''
    ROS Subscriber
    '''
    self.waypoint_sub = rospy.Subscriber(f"{self.drone_type}_{self.drone_id}/move_base_simple/goal", 
                                     PoseStamped, self.waypoint_callback)
    self.local_pose_sub = rospy.Subscriber(f"{self.drone_type}_{str(self.drone_id)}/mavros/local_position/pose",
                                           PoseStamped, self.local_pose_callback)
    self.other_drone_pose_sub = [None] * self.drone_num
    self.other_vehicle_pose_sub = [None] * self.vehicle_num
    '''
    ROS Publisher
    '''
    self.vel_flu_pub = rospy.Publisher(f"{self.drone_type}_{str(self.drone_id)}/cmd_vel_flu", Twist, queue_size=1)
    self.cmd_pub = rospy.Publisher(f"{self.drone_type}_{str(self.drone_id)}/cmd", String, queue_size=10)
  
  def waypoint_callback(self, msg):
    print("trigger!!!")
    self.target_position = msg.pose.position
    x = msg.pose.orientation.x
    y = msg.pose.orientation.y
    z = msg.pose.orientation.z
    w = msg.pose.orientation.w
    rot = PyKDL.Rotation.Quaternion(x, y, z, w)
    
    res = rot.GetRPY()[2]
    while res > math.pi:
      res -= 2.0 * math.pi
    while res < -math.pi:
      res += 2.0 * math.pi
    if res < 0:
      res = res  + 2.0 * math.pi
    self.target_yaw = res
  
  def local_pose_callback(self, msg):
    self.current_pose = msg
    self.current_position = self.current_pose.pose.position
    self.current_position.x = self.current_position.x + self.bias[self.drone_id][0]
    self.current_position.y = self.current_position.y + self.bias[self.drone_id][1]
    self.current_position.z = self.current_position.z
    x = self.current_pose.pose.orientation.x
    y = self.current_pose.pose.orientation.y
    z = self.current_pose.pose.orientation.z
    w = self.current_pose.pose.orientation.w
    rot = PyKDL.Rotation.Quaternion(x, y, z, w)
    
    res = rot.GetRPY()[2]
    while res > math.pi:
      res -= 2.0 * math.pi
    while res < -math.pi:
      res += 2.0 * math.pi
    if res < 0:
      res = res  + 2.0 * math.pi
    self.current_yaw = res   # [0, 2pi]
  
  def other_drone_pose_callback(self, msg, id):
    self.other_drone_pose[id] = msg
    self.other_drone_pose[id].pose.position.x = self.other_drone_pose[id].pose.position.x + self.bias[id][0]
    self.other_drone_pose[id].pose.position.y = self.other_drone_pose[id].pose.position.y + self.bias[id][1]
    self.other_drone_pose[id].pose.position.z = self.other_drone_pose[id].pose.position.z
  
  def other_vehicle_pose_callback(self, msg, id):
    self.other_vehicle_pose[id] = msg
    
  def run(self):
    rate = rospy.Rate(self.rate)
    for i in range(self.drone_num):
      if not i == self.drone_id:
        self.other_drone_pose_sub[i] = rospy.Subscriber(f"{self.drone_type}_{str(i)}/mavros/local_position/pose",
                                                        PoseStamped, self.other_drone_pose_callback, i)
    for i in range(self.vehicle_num):
      self.other_vehicle_pose_sub[i] = rospy.Subscriber(f"vehicle_{str(i)}/global_position/pose",
                                                        PoseStamped, self.other_vehicle_pose_callback, i)
    while not rospy.is_shutdown():
      self.count += 1
      self.cmd = ''
      if self.exec_state == 0:
        if self.count == 40 or self.count == 42 or self.count == 44:
          self.cmd = 'OFFBOARD'
        if self.count == 94 or self.count == 96 or self.count == 98:
          self.cmd = 'ARM'
        self.target_position.linear.z = self.height_limit
        self.target_position.linear.x = self.current_position.x
        self.target_position.linear.y = self.current_position.y
        self.target_motion.angular.x = self.target_motion.angular.x
        self.target_motion.angular.y = self.target_motion.angular.y
        self.target_motion.angular.z = self.target_motion.angular.z
      
      elif self.exec_state == 1 and self.change_task_flag:
        self.change_task_flag = False
        self.enable_receive_goal = True
        self.target_motion.angular.x = 0
        self.target_motion.angular.y = 0
        self.target_motion.angular.z = 0
              
      distance_tar_cur = self.VectNorm3(self.target_position.linear, self.current_position)
      if distance_tar_cur < 1:
        self.arrive_count += 1
        if self.arrive_count > 3:
          self.is_arrive = True
          self.arrive_count = 0
        else:
          self.is_arrive = False
      else:
        self.arrive_count = 0
        self.is_arrive = False
        
      # task change
      if (self.exec_state == 0) and self.is_arrive:
        self.change_task_flag = True
        self.exec_state = 1
        self.is_arrive = False
      
      self.get_control_vel()
      self.drone_avoid()
      self.vel_flu_pub.publish(self.target_motion)
      self.cmd_pub.publish(self.cmd)
      
      rate.sleep()
    
  def get_control_vel(self): # PID control, flu coordinate
    dis_curtar = self.VectNorm2(self.target_position.linear, self.current_position)
    vel_total = self.kp * dis_curtar  # velocity
    if vel_total > self.vxy_max:
        vel_total = self.vxy_max
    self.target_yaw = self.pos2ang(self.target_position.linear.x, self.target_position.linear.y, self.current_position.x, self.current_position.y)
    mid_yaw = self.target_yaw - self.current_yaw 
    if mid_yaw > math.pi:
        mid_yaw = mid_yaw - 2 * math.pi
    elif mid_yaw < -math.pi:
        mid_yaw = 2 * math.pi + mid_yaw
    self.target_motion.angular.z = self.ky * mid_yaw
    if self.target_motion.angular.z > self.angz_max:
        self.target_motion.angular.z = self.angz_max
    elif self.target_motion.angular.z < -self.angz_max:
        self.target_motion.angular.z = -self.angz_max
    self.target_motion.linear.x  = vel_total * math.cos(mid_yaw)
    self.target_motion.linear.y = vel_total * math.sin(mid_yaw)

    self.target_motion.linear.z = self.kp * (self.target_position.linear.z - self.current_position.z)
  
  def drone_avoid(self):
    self.avoid = []
    for i in range(self.drone_num):
        if not i == self.drone_id:
            dis_partner = math.sqrt(
                (self.current_position.x - self.other_drone_pose[i].pose.position.x) ** 2 + (self.current_position.y - self.other_drone_pose[i].pose.position.y) ** 2
            )
            if dis_partner < self.safe_distance:
                if (self.current_position.z - self.other_drone_pose[i].pose.position.z) < self.safe_height:
                  self.avoid.append(i)
    avoid_num = len(self.avoid)
    heigher_num = 0
    if avoid_num > 0:
        for j in range(avoid_num):
            if self.other_drone_pose[self.avoid[j]].pose.position.z > self.current_position.z:
                heigher_num = heigher_num + 1
        if heigher_num == 0:  
            self.target_position.linear.z = self.height_limit + self.safe_height
        else:                 
            self.target_position.linear.z = self.height_limit - self.safe_height * heigher_num
    else:
        self.target_position.linear.z = self.height_limit
  
  def pso_algorithm(self):
    if self.enable_receive_goal:
      # TODO:
      self.target_motion.linear.x = 0.0
      self.target_motion.linear.y = 0.0
  
  def VectNorm3(self, pt1, pt2):
    norm = math.sqrt(pow(pt1.x - pt2.x, 2) + pow(pt1.y - pt2.y, 2) + pow(pt1.z - pt2.z, 2))
    return norm
  
  def VectNorm2(self, pt1, pt2):
    norm = math.sqrt(pow(pt1.x - pt2.x, 2) + pow(pt1.y - pt2.y, 2))
    return norm
  
  def pos2ang(self, xa, ya, xb, yb):
    if not xa-xb == 0:
      angle = math.atan2((ya - yb),(xa - xb))
      if (ya-yb > 0) and (angle < 0):  
          angle = angle + math.pi       
      elif (ya-yb < 0) and (angle > 0): 
          angle = angle - math.pi       
      elif ya-yb == 0:
          if xa-xb > 0:
              angle = 0.0
          else:
              angle = math.pi
    else:
      if ya-yb > 0:
          angle = math.pi / 2
      elif ya-yb <0:
          angle = -math.pi / 2
      else:
          angle = 0.0
    if angle < 0:
      angle = angle + 2 * math.pi   # 0 to 2pi
    return angle
  
if __name__ == '__main__':
  planner = NavigationPlan(sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4])
  planner.run()